"""
矩阵运算模块
包含所有矩阵相关运算函数
"""
import numpy as np
from typing import List, Union

def matrix_add(a: List[List[Union[int, float]]], b: List[List[Union[int, float]]]) -> List[List[Union[int, float]]]:
    """
    矩阵加法
    :param a: 第一个矩阵
    :param b: 第二个矩阵
    :return: 两个矩阵的和
    """
    try:
        a = np.array(a)
        b = np.array(b)
        if a.shape != b.shape:
            raise ValueError("Matrices must have the same shape")
        return (a + b).tolist()
    except Exception as e:
        raise ValueError(f"Matrix addition error: {str(e)}")

def matrix_multiply(a: List[List[Union[int, float]]], b: List[List[Union[int, float]]]) -> List[List[Union[int, float]]]:
    """
    矩阵乘法
    :param a: 第一个矩阵
    :param b: 第二个矩阵
    :return: 两个矩阵的乘积
    """
    try:
        a = np.array(a)
        b = np.array(b)
        if a.shape[1] != b.shape[0]:
            raise ValueError("Number of columns in first matrix must equal number of rows in second")
        return np.matmul(a, b).tolist()
    except Exception as e:
        raise ValueError(f"Matrix multiplication error: {str(e)}")

def matrix_transpose(a: List[List[Union[int, float]]]) -> List[List[Union[int, float]]]:
    """
    矩阵转置
    :param a: 输入矩阵
    :return: 转置后的矩阵
    """
    try:
        return np.array(a).T.tolist()
    except Exception as e:
        raise ValueError(f"Matrix transpose error: {str(e)}")

def matrix_determinant(a: List[List[Union[int, float]]]) -> float:
    """
    计算矩阵行列式
    :param a: 方阵
    :return: 行列式值
    """
    try:
        a = np.array(a)
        if a.shape[0] != a.shape[1]:
            raise ValueError("Matrix must be square")
        return float(np.linalg.det(a))
    except Exception as e:
        raise ValueError(f"Matrix determinant error: {str(e)}")

def matrix_inverse(a: List[List[Union[int, float]]]) -> List[List[Union[int, float]]]:
    """
    计算矩阵逆
    :param a: 可逆方阵
    :return: 逆矩阵
    """
    try:
        a = np.array(a)
        if a.shape[0] != a.shape[1]:
            raise ValueError("Matrix must be square")
        return np.linalg.inv(a).tolist()
    except Exception as e:
        raise ValueError(f"Matrix inverse error: {str(e)}")

def matrix_eigen(a: List[List[Union[int, float]]]) -> dict:
    """
    计算矩阵的特征值和特征向量
    :param a: 方阵
    :return: 字典包含特征值(eigenvalues)和特征向量(eigenvectors)
    """
    try:
        a = np.array(a)
        if a.shape[0] != a.shape[1]:
            raise ValueError("Matrix must be square")
        eigenvalues, eigenvectors = np.linalg.eig(a)
        return {
            "eigenvalues": eigenvalues.tolist(),
            "eigenvectors": eigenvectors.tolist()
        }
    except Exception as e:
        raise ValueError(f"Matrix eigen calculation error: {str(e)}")

def characteristic_polynomial(a: List[List[Union[int, float]]]) -> List[float]:
    """
    计算矩阵的特征多项式系数
    :param a: 方阵
    :return: 特征多项式系数列表，从高次到低次
    """
    try:
        a = np.array(a)
        if a.shape[0] != a.shape[1]:
            raise ValueError("Matrix must be square")
        # 计算特征多项式系数
        char_poly = np.poly(a)
        return char_poly.tolist()
    except Exception as e:
        raise ValueError(f"Characteristic polynomial calculation error: {str(e)}")
        
def minimal_polynomial(a: List[List[Union[int, float]]]) -> List[float]:
    """
    计算矩阵的极小多项式系数
    :param a: 方阵
    :return: 极小多项式系数列表，从高次到低次
    """
    try:
        a = np.array(a)
        if a.shape[0] != a.shape[1]:
            raise ValueError("Matrix must be square")
            
        # 计算特征值
        eigenvalues = np.linalg.eigvals(a)
        unique_eigenvalues = np.unique(eigenvalues)
        
        # 计算每个特征值的代数重数
        min_poly = np.poly1d([1])
        for ev in unique_eigenvalues:
            # 计算该特征值的Jordan块最大阶数
            rank_changes = []
            mat = a - ev * np.eye(a.shape[0])
            for k in range(1, a.shape[0]+1):
                rank = np.linalg.matrix_rank(mat**k)
                if k > 1:
                    rank_changes.append(prev_rank - rank)
                prev_rank = rank
            
            # 确定该特征值对应的极小多项式因子次数
            max_jordan = max(rank_changes) if rank_changes else 1
            min_poly *= np.poly1d([1, -ev], r=False)**max_jordan
            
        return min_poly.coeffs.tolist()
    except Exception as e:
        raise ValueError(f"Minimal polynomial calculation error: {str(e)}")